{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from bertviz.bertviz.attention_details_gpt2 import AttentionDetailsGpt2Data, AttentionDetailsGptData, show\n",
    "from bertviz.bertviz.pytorch_pretrained_bert import OpenAIGPTLMHeadModel \n",
    "from pytorch_pretrained_bert import OpenAIGPTTokenizer \n",
    "import os\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "require.config({\n",
    "  paths: {\n",
    "      d3: '//cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min'\n",
    "  }\n",
    "});"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_html():\n",
    "  import IPython\n",
    "  display(IPython.core.display.HTML('''\n",
    "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
    "        <script>\n",
    "          requirejs.config({\n",
    "            paths: {\n",
    "              base: '/static/base',\n",
    "              \"d3\": \"https://cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min\",\n",
    "              jquery: '//ajax.googleapis.com/ajax/libs/jquery/3.3.1/jquery.min',\n",
    "            },\n",
    "          });\n",
    "        </script>\n",
    "        '''))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "special_tokens = ['<ATTR_WORDS>','<CON_START>','<START>','<END>']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = OpenAIGPTLMHeadModel.from_pretrained('openai-gpt',num_special_tokens=len(special_tokens))\n",
    "tokenizer = OpenAIGPTTokenizer.from_pretrained('openai-gpt', special_tokens = special_tokens)\n",
    "#path =  \"/home/ubuntu/bhargav/log_yelp_retireve_edit_bert_best_head_preprocessing/pytorch_model_zero_grad_1.bin\"\n",
    "#model_state_dict = torch.load(path)\n",
    "#model.load_state_dict(model_state_dict)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = GPT2Model.from_pretrained('gpt2')\n",
    "#tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "\n",
    "text = \"<ATTR_WORDS> awesome delicious <CON_START> service here is and food is <START> service here is awesome and food is \"\n",
    "#details_data = AttentionDetailsGpt2Data(model, tokenizer)\n",
    "details_data = AttentionDetailsGptData(model, tokenizer)\n",
    "tokens, queries, keys, atts = details_data.get_data(text)\n",
    "call_html()\n",
    "show(tokens, queries, keys, atts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = GPT2Model.from_pretrained('gpt2')\n",
    "#tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "\n",
    "text = \"<ATTR_WORDS> great <CON_START> this place is ! <START> this place is \"\n",
    "#details_data = AttentionDetailsGpt2Data(model, tokenizer)\n",
    "details_data = AttentionDetailsGptData(model, tokenizer)\n",
    "tokens, queries, keys, atts = details_data.get_data(text)\n",
    "call_html()\n",
    "show(tokens, queries, keys, atts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
